import uuid
import datetime
from datetime import timezone
from functools import wraps
from typing import Callable, Any, Optional
import redis
from redis.exceptions import ConnectionError

from flask_jwt_extended import get_jwt_identity, verify_jwt_in_request, create_access_token
from flask import request, jsonify
from models import User, db
from models.auth_log import AuthLog
from flask import current_app


import uuid
from logging_config import logger

# 权限常量定义（与模型完全一致）
ROLE_USER = 'user'
ROLE_ADMIN = 'admin'
ROLE_MODERATOR = 'moderator'
ROLE_LAW = 'law_enforcement'

# 初始化Redis连接（带连接池和异常处理）
def get_redis_client(app):
    import os
    """获取Redis客户端（带连接池和重连机制）"""
    # 开发环境下直接使用内存模拟的Redis
    if app.config.get('ENV', 'production') == 'development':
        app.logger.info("开发环境下使用内存模拟Redis客户端")
        from utils.memory_redis import MemoryRedis
        return MemoryRedis()
    
    # 生产环境尝试连接Redis服务器
    try:
        # 优先从环境变量获取REDIS_URL
        redis_url = os.getenv('REDIS_URL')
        
        if redis_url:
            # 使用from_url方法连接Redis
            client = redis.Redis.from_url(
                redis_url,
                decode_responses=True,
                socket_connect_timeout=5
            )
            # 测试连接
            client.ping()
            app.logger.info(f"成功连接到Redis: {redis_url}")
            return client
        else:
            # 如果没有REDIS_URL，则使用单独的配置项
            redis_host = os.getenv('REDIS_HOST', 'localhost')
            redis_port = int(os.getenv('REDIS_PORT', 6379))
            redis_db = int(os.getenv('REDIS_DB', 0))
            redis_password = os.getenv('REDIS_PASSWORD', '')

            client = redis.Redis(
                host=redis_host,
                port=redis_port,
                db=redis_db,
                password=redis_password,
                decode_responses=True,
                socket_connect_timeout=5
            )
            # 测试连接
            client.ping()
            app.logger.info(f"成功连接到Redis: {redis_host}:{redis_port}/{redis_db}")
            return client
    except Exception as e:
        app.logger.error(f"Redis连接失败: {str(e)}")
        app.logger.warning("将回退到使用内存模拟的Redis")
        from utils.memory_redis import MemoryRedis
        return MemoryRedis()



def get_current_user() -> Optional[User]:
    """
    获取当前登录用户对象（带缓存优化和异常处理）
    
    :return: User对象或None
    """
    try:
        user_id = get_jwt_identity()
        if not user_id:
            return None
            
        user_id_uuid = uuid.UUID(user_id)
        
        # 获取Redis客户端
        redis_client = get_redis_client(current_app)
        
        # 尝试从缓存获取用户
        cache_key = f"user:{user_id}"
        if redis_client:
            cached_user_id = redis_client.get(cache_key)
            if cached_user_id and uuid.UUID(cached_user_id) == user_id_uuid:
                return User.query.get(user_id_uuid)
        
        # 从数据库查询并缓存
        user = User.query.get(user_id_uuid)
        if user and redis_client:
            redis_client.setex(cache_key, 300, str(user_id_uuid))  # 缓存5分钟
            
        return user
    except (ValueError, TypeError) as e:
        logger.error(f"JWT中用户ID格式错误: {str(e)}")
        return None
    except Exception as e:
        logger.error(f"获取当前用户异常: {str(e)}", exc_info=True)
        return None

def require_authentication(f: Callable) -> Callable:
    """
    用户认证装饰器（带双因素认证和账号锁定检查）
    
    :param f: 被装饰的函数
    :return: 装饰后的函数
    """
    @wraps(f)
    def decorated_function(*args: Any, **kwargs: Any) -> Any:
        verify_jwt_in_request()
        user = get_current_user()
        
        if not user:
            log_auth_event(None, 'JWT_INVALID', request.remote_addr)
            return jsonify({'error': '认证失败，请重新登录'}), 401
        
        # 检查账号是否被锁定（使用带时区的时间）
        now = datetime.datetime.now(timezone.utc)
        if user.login_locked_until:
            # 确保login_locked_until是时区感知的
            locked_until = user.login_locked_until
            if locked_until.tzinfo is None:
                # 如果是naive datetime，假设为UTC
                locked_until = locked_until.replace(tzinfo=timezone.utc)
            
            if locked_until > now:
                lock_seconds = (locked_until - now).total_seconds()
                log_auth_event(user.id, 'ACCOUNT_LOCKED', request.remote_addr)
                return jsonify({
                    'error': '账号已锁定', 
                    'lock_seconds': int(lock_seconds)
                }), 423
        
        # 检查账号是否活跃
        if not user.is_active:
            log_auth_event(user.id, 'ACCOUNT_DISABLED', request.remote_addr)
            return jsonify({'error': '账号已禁用'}), 403
        
        # 检查双因素认证（如已启用）
        # 注意：two_factor_enabled字段在当前User模型中不存在
        # 如果需要双因素认证，请在User模型中添加此字段
        # if hasattr(user, 'two_factor_enabled') and user.two_factor_enabled:
        #     totp_token = request.headers.get('X-TOTP-Token')
        #     if not totp_token or not user.verify_totp(totp_token):
        #         log_auth_event(user.id, '2FA_FAILED', request.remote_addr)
        #         return jsonify({'error': '双因素认证失败'}), 401
        
        return f(user, *args, **kwargs)
    return decorated_function

def require_role(required_role: str) -> Callable:
    """
    角色权限验证装饰器（记录详细权限日志）
    
    :param required_role: 需要的角色
    :return: 装饰器函数
    """
    def decorator(f: Callable) -> Callable:
        @wraps(f)
        def decorated_function(user: User, *args: Any, **kwargs: Any) -> Any:
            if user.user_role != required_role:
                logger.warning(f"权限拒绝: 用户ID={user.id} 角色={user.user_role} "
                              f"需要角色={required_role} 访问端点={request.path}")
                log_auth_event(user.id, 'PERMISSION_DENIED', request.remote_addr)
                return jsonify({'error': '权限不足'}), 403
            return f(user, *args, **kwargs)
        return decorated_function
    return decorator

def generate_jwt(user_id: str, expires_hours: int = 1) -> str:
    """
    生成JWT令牌（带自定义过期时间和完整声明）
    
    :param user_id: 用户ID
    :param expires_hours: 令牌过期小时数
    :return: JWT令牌字符串
    """
    try:
        user_id_uuid = uuid.UUID(user_id)
        user = User.query.get(user_id_uuid)
        if not user:
            raise ValueError("用户不存在")
            
        expires = datetime.timedelta(hours=expires_hours)
        return create_access_token(
            identity=str(user_id_uuid),
            expires_delta=expires,
            additional_claims={
                'role': user.user_role,
                'user_name': user.username,
                'iat': datetime.datetime.now(timezone.utc).isoformat(),
                'jti': str(uuid.uuid4())  # 唯一令牌ID
            }
        )
    except Exception as e:
        logger.error(f"生成JWT令牌异常: {str(e)}", exc_info=True)
        raise

def validate_jwt(token: str) -> Optional[dict]:
    """
    验证JWT令牌有效性（返回完整载荷）
    
    :param token: JWT令牌
    :return: 令牌载荷或None
    """
    from flask_jwt_extended import decode_token, JWTError
    try:
        decoded_token = decode_token(token)
        # 验证用户是否存在且活跃
        if not is_user_active(decoded_token['sub']):
            logger.warning(f"JWT令牌用户已失效: {decoded_token['sub']}")
            return None
        return decoded_token
    except JWTError as e:
        logger.warning(f"JWT格式错误: {str(e)}")
    except Exception as e:
        logger.error(f"JWT验证异常: {str(e)}", exc_info=True)
    return None

def is_user_active(user_id: str) -> bool:
    """
    检查用户是否活跃（优化查询）
    
    :param user_id: 用户ID
    :return: 活跃状态
    """
    try:
        user_id_uuid = uuid.UUID(user_id)
        user = User.query.filter_by(
            id=user_id_uuid, 
            is_active=True
        ).first()
        return bool(user)
    except Exception as e:
        logger.error(f"检查用户活跃状态异常: {str(e)}")
        return False

import datetime
from datetime import timezone
from typing import Optional
from models.auth_log import AuthLog
from models import db

def log_auth_event(user_id: Optional[str], event_type: str, ip_address: str, user_agent: str = None) -> None:
    """
    记录认证事件日志（增强健壮性）
    
    :param user_id: 用户ID
    :param event_type: 事件类型
    :param ip_address: IP地址
    :param user_agent: 用户代理
    """
    try:
        log = AuthLog(
            user_id=user_id,
            event_type=event_type,
            ip_address=ip_address,
            user_agent=user_agent or (request.user_agent.string if request else "unknown"),
            timestamp=datetime.datetime.now(timezone.utc)
        )
        db.session.add(log)
        db.session.commit()
        logger.info(f"认证事件: {event_type} 用户ID={user_id} IP={ip_address}")
    except Exception as e:
        logger.error(f"记录认证日志异常: {str(e)}", exc_info=True)
        # 确保异常不影响主流程
        db.session.rollback()

def track_login_attempt(user_id: Optional[str], ip_address: str, success: bool = False) -> bool:
    """
    跟踪登录尝试，实现频率限制（增强安全性）
    
    :param user_id: 用户ID
    :param ip_address: IP地址
    :param success: 登录是否成功
    :return: 登录是否允许
    """
    # 获取Redis客户端
    redis_client = get_redis_client(current_app)
    
    user = None
    try:
        if user_id:
            user_id_uuid = uuid.UUID(user_id)
            user = User.query.get(user_id_uuid)
    except ValueError:
        logger.warning(f"无效的用户ID: {user_id}")
    
    # 记录登录尝试
    event_type = 'LOGIN_SUCCESS' if success else 'LOGIN_FAILED'
    log_auth_event(user_id, event_type, ip_address, request.user_agent.string if request else None)
    
    if success:
        # 登录成功，重置尝试次数
        if user:
            user.failed_login_attempts = 0
            user.login_locked_until = None
        # 刷新用户缓存
        if redis_client and user_id:
            redis_client.delete(f"user:{user_id}")
        return True
    
    # 登录失败，处理尝试次数 - 暂时禁用频率限制
    if user:
        user.failed_login_attempts += 1
        # 频率限制已禁用
        pass  # 不做任何限制
    else:
        # 匿名用户限制已禁用
        pass  # 不做任何限制
    
    db.session.commit()
    return True

def revoke_jwt(token: str) -> bool:
    """
    将JWT加入黑名单（支持令牌提前失效）
    
    :param token: JWT令牌
    :return: 是否成功
    """
    try:
        decoded = validate_jwt(token)
        if not decoded:
            return False
            
        jti = decoded.get('jti')
        if not jti:
            return False
            
        # 获取令牌剩余过期时间
        exp = decoded.get('exp', 0)
        now = int(datetime.datetime.now(timezone.utc).timestamp())
        remaining = exp - now
        
        if remaining > 0 and redis_client:
            redis_client.setex(f"jwt_blacklist:{jti}", int(remaining), "revoked")
            logger.info(f"JWT已加入黑名单: jti={jti} 剩余有效期={remaining}秒")
            return True
        return False
    except Exception as e:
        logger.error(f"revoke_jwt异常: {str(e)}", exc_info=True)
        return False

def is_jwt_revoked(token: str) -> bool:
    """
    检查JWT是否已被吊销
    
    :param token: JWT令牌
    :return: 是否已吊销
    """
    try:
        decoded = validate_jwt(token)
        if not decoded:
            return True  # 无效令牌视为已吊销
            
        jti = decoded.get('jti')
        if not jti or not redis_client:
            return False
            
        return bool(redis_client.get(f"jwt_blacklist:{jti}"))
    except Exception as e:
        logger.error(f"is_jwt_revoked异常: {str(e)}", exc_info=True)
        return True